import matplotlib.pyplot as plt
import numpy as np
import sqlite3
from datetime import datetime

# 定义函数，用于处理数据
def process_data(rows):
    dates = [datetime.strptime(row[0], '%Y-%m-%d') for row in rows]
    total_market_values = [float(row[1]) if row[1] else None for row in rows]
    circulating_market_values = [float(row[2]) if row[2] else None for row in rows]
    listed_companies = [int(row[3]) if row[3] else None for row in rows]
    average_pe_ratios = [float(row[4]) if row[4] else None for row in rows]

    # 计算流通市值占总市值的比例
    ratios = []
    for total, circulating in zip(total_market_values, circulating_market_values):
        if total and circulating:
            ratio = float(circulating) / float(total)
            ratios.append(ratio)
        else:
            ratios.append(None)

    # 过滤掉 None 值
    valid_dates = []
    valid_total_market_values = []
    valid_circulating_market_values = []
    valid_ratios = []
    for j in range(len(dates)):
        if total_market_values[j] and circulating_market_values[j] and ratios[j]:
            valid_dates.append(dates[j])
            valid_total_market_values.append(total_market_values[j])
            valid_circulating_market_values.append(circulating_market_values[j])
            valid_ratios.append(ratios[j])

    return valid_dates, valid_total_market_values, valid_circulating_market_values, valid_ratios

# 定义函数，用于绘制子图
def plot_subplot(ax, dates, values, label, title, xlabel, ylabel):
    if len(dates) == 1:
        ax.scatter(dates, values, label=label)
    else:
        # 绘制曲线
        ax.plot(dates, values, label=label)
        # 绘制数据点
        ax.scatter(dates, values, color='red', zorder=5)  # zorder=5 确保数据点显示在曲线之上
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.legend()

# 定义函数，用于绘制数据变化趋势图
def plot_data_trend():
    try:
        # 连接到数据库
        conn = sqlite3.connect('stock_data.db')
        cursor = conn.cursor()

        # 指定支持中文的字体，这里以 SimHei 为例，你可以根据自己系统中的字体进行调整
        plt.rcParams['font.family'] = 'SimHei'
        # 解决负号显示问题
        plt.rcParams['axes.unicode_minus'] = False

        exchanges = ['上海交易所', '深圳交易所']
        plt.figure(figsize=(16, 12))

        for i, exchange in enumerate(exchanges):
            cursor.execute('SELECT date, total_market_value, circulating_market_value, listed_companies, average_pe_ratio FROM stock_data WHERE exchange =?', (exchange,))
            rows = cursor.fetchall()

            if not rows:
                print(f"未找到 {exchange} 的数据")
                continue

            # 打印处理前的数据内容
            print(f"处理前 {exchange} 的数据内容:")
            for row in rows:
                print(row)

            valid_dates, valid_total_market_values, valid_circulating_market_values, valid_ratios = process_data(rows)

            # 打印处理后的数据内容
            print(f"处理后 {exchange} 的有效日期数据: {valid_dates}")
            print(f"处理后 {exchange} 的有效总市值数据: {valid_total_market_values}")
            print(f"处理后 {exchange} 的有效流通市值数据: {valid_circulating_market_values}")
            print(f"处理后 {exchange} 的有效比例数据: {valid_ratios}")

            if not valid_dates:
                print(f"{exchange} 没有有效的数据用于绘图")
                continue

            # 绘制子图
            ax1 = plt.subplot(2, 2, i * 2 + 1)
            plot_subplot(ax1, valid_dates, valid_total_market_values, 'Total Market Value', f'{exchange} Market Value Trend', 'Date', 'Market Value (Billion Yuan)')
            plot_subplot(ax1, valid_dates, valid_circulating_market_values, 'Circulating Market Value', f'{exchange} Market Value Trend', 'Date', 'Market Value (Billion Yuan)')

            ax2 = plt.subplot(2, 2, i * 2 + 2)
            plot_subplot(ax2, valid_dates, valid_ratios, 'Circulating / Total Market Value Ratio', f'{exchange} Circulating / Total Market Value Ratio Trend', 'Date', 'Ratio')

        plt.tight_layout()
        plt.show()

    except sqlite3.Error as e:
        print(f"数据库错误: {e}")
    finally:
        # 关闭数据库连接
        if conn:
            conn.close()

if __name__ == "__main__":
    plot_data_trend()